package com.vshop.framework.web.trace;

import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.ttl.TransmittableThreadLocal;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.concurrent.atomic.AtomicInteger;

public class MdcTraceUtil {

    /**
     * 追踪id的名称
     */
    public static final String TraceIdKey = "traceId";
    /**
     * 块id的名称
     */
    public static final String SpanIdKey = "spanId";

    /**
     * 日志链路追踪id信息头
     */
    public static final String TraceIdHeaderKey = "x-traceId-header";

    /**
     * 日志链路块id信息头
     */
    public static final String SpanIdHeaderKey = "x-spanId-header";


    private static final TransmittableThreadLocal<AtomicInteger> spanNumber = new TransmittableThreadLocal<>();
    private static final Logger log = LoggerFactory.getLogger(MdcTraceUtil.class);

    /**
     * 创建traceId并赋值MDC
     */
    public static void addTrace() {
        String traceId = createTraceId();
        MDC.put(TraceIdKey, traceId);
        MDC.put(SpanIdKey, "0");
        initSpanNumber();
    }

    /**
     * 赋值MDC
     */
    public static void putTrace(String traceId, String spanId) {
        MDC.put(TraceIdKey, traceId);
        MDC.put(SpanIdKey, spanId);
        initSpanNumber();
    }

    /**
     * 获取MDC中的traceId值
     */
    public static String getTraceId() {
        return MDC.get(TraceIdKey);
    }

    /**
     * 获取MDC中的spanId值
     */
    public static String getSpanId() {
        return MDC.get(SpanIdKey);
    }

    /**
     * 清除MDC的值
     */
    public static void removeTrace() {
        MDC.remove(TraceIdKey);
        MDC.remove(SpanIdKey);
        spanNumber.remove();
    }

    /**
     * 创建traceId
     */
    public static String createTraceId() {
        return IdUtil.getSnowflake().nextIdStr();
    }

    public static String getNextSpanId() {
        return StrUtil.format("{}.{}", getSpanId(), spanNumber.get().incrementAndGet());
    }

    private static void initSpanNumber() {
        spanNumber.set(new AtomicInteger(0));
    }


}
